import math
from functools import partial
from itertools import repeat

import torch
import torch.nn as nn
import torch.nn.functional as F
# from torch._six import container_abcs
import collections.abc as container_abcs
int_classes = int
string_classes = str

import torch.utils.checkpoint as checkpoint
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
# from timm.models.vision_transformer import Mlp as Mlp_old
# from layers.ns_structure_loss import Attention
from .resnet import ResNet, BasicBlock, Bottleneck
# from modeling.fusion_mo import ChannelAttentionModule    # , ASFFS
import copy
import random
from torch.nn import Parameter
# from .gcn import GraphConvolution, gen_adj_sim, norm_adj_batch    # , gen_adj_sim_x
# from .gcn import HRCNHead
SPATIAL_AND_CHANNEL = True

def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')
        if m.bias is not None:
          nn.init.constant_(m.bias, 0.0)
    elif classname.find('Conv') != -1:
        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
        if m.bias is not None:
            nn.init.constant_(m.bias, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine:
            nn.init.constant_(m.weight, 1.0)
            nn.init.constant_(m.bias, 0.0)


class LayerScale_Block_CA(nn.Module): 
    # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
    # with slight modifications to add CA and LayerScale
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, Attention_block = 'Class_Attention',
                 Mlp_block='Mlp',init_values=1e-4,num_stripes=1):
        super().__init__()
        # self.attn = Attention_block(
        #     dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.q = nn.Linear(2048, 2048, bias=False)
        self.k = nn.Linear(2048, 2048, bias=False)        
        self.q2 = nn.Linear(512, 512, bias=False)
        self.k2 = nn.Linear(512, 512, bias=False)        
        # self.relu = nn.LeakyReLU(0.2)
        self.relu = nn.ReLU()
        
        self.gcn1g = GraphConvolution(2048, 2048)
        self.gcn2g = GraphConvolution(2048, 2048)
        # self.attn2 = Attention_block(
        #     int(dim/num_stripes), num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.gcnl_1l = nn.ModuleList()
        self.gcnl_2l = nn.ModuleList()
        for i in range(4):
            self.gcnl_1l.append(GraphConvolution(512, 512, win_weight=512, wout_weight=512))
            self.gcnl_2l.append(GraphConvolution(512, 512, win_weight=512, wout_weight=512))
                
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.num_stripes=num_stripes
        mlp_hidden_dim = int(dim * mlp_ratio)
        # self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        # self.mlp2 = Mlp_block(in_features=int(dim/num_stripes), hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.norm_list=nn.ModuleList()
        self.gamma_list=[]
        for i in range(2):
            self.gamma = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True).cuda()
            self.norm_list.append(norm_layer(dim))
            self.gamma_list.append(self.gamma)
        for i in range(num_stripes*2):
            self.gamma = nn.Parameter(init_values * torch.ones((int(dim/num_stripes))),requires_grad=True).cuda()
            self.norm_list.append(norm_layer(int(dim/num_stripes)))
            self.gamma_list.append(self.gamma)
    
    def forward(self, global_feat,local_feats, cls_tokens):
        B = global_feat.shape[0]
        # print(global_feat.shape)
        # print(cls_tokens[0][:B].shape)
        # raise "16516063"
        u1 = torch.cat((cls_tokens[0][:B], global_feat),dim=1)    # B, 197. 2048
        u_feats=[u1]
        # print(u1.shape)
        for i in range(len(local_feats)):
            # raise "5111515155"
            u_tmp=torch.cat((cls_tokens[i+1][:B], local_feats[i].view(B, 512, -1).transpose(1,2).contiguous()),dim=1)    # B, 197. 512
            # print(u_tmp.shape)
            # raise "ssss"
            # torch.Size([64, 257, 2048])   torch.Size([64, 257, 512])
            u_feats.append(u_tmp) 
        
        gfeature = u1    # self.q(
        edge_feature = u1    # self.k
        # print(edge_feature.shape)
        edge = gen_adj_sim(edge_feature, edge_feature.permute(0, 2, 1))
        
        # conduct edge normalization
        adj = edge.cuda()

        adj = norm_adj_batch(adj)            
        gcn_r = self.relu(self.gcn1g(gfeature, adj))
        gcn_rg = self.relu(self.gcn2g(gcn_r, adj))

        x_cls_list = [gcn_rg]
       
        for i in range(self.num_stripes):
            lfeature = u_feats[i+1]    # self.q(
            edge_lfeature = u_feats[i+1]    # self.k
            # print(edge_feature.shape)
            edgel = gen_adj_sim(edge_lfeature, edge_lfeature.permute(0, 2, 1))
            # conduct edge normalization
            adjl = edgel.cuda()
            adjl = norm_adj_batch(adjl)    
            gcn_lr = self.relu(self.gcnl_1l[i](lfeature, adjl))
            gcn_rl = self.relu(self.gcnl_2l[i](gcn_lr, adjl))
            x_cls_list.append(gcn_rl)

        # x_cls1 = cls_tokens[0] + self.drop_path(self.gamma_list[0] * self.attn(self.norm_list[0](u1)))
        # # print(self.gamma_list[0] * self.attn(self.norm_list[0](u1)))
        # x_cls1 = x_cls1 + self.drop_path(self.gamma_list[1] * self.mlp(self.norm_list[1](x_cls1)))

        # x_cls_list=[x_cls1]
        # for i in range(self.num_stripes):
        #     x_cls2 = cls_tokens[i+1] + self.drop_path(self.gamma_list[2+i*2] * self.attn2(self.norm_list[2+i*2](u_feats[i+1])))

        #     x_cls2 = x_cls2 + self.drop_path(self.gamma_list[3+i*2] * self.mlp2(self.norm_list[3+i*2](x_cls2)))
        #     x_cls_list.append(x_cls2)
        return x_cls_list

class GraphConvolution(nn.Module):
    """
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907

    in_features: input feature dimension
    out_features: output feature dimension

    dropout: default=0 not used， call nn.dropout function

    bias: the learnable bias in GCN layer, default set as True

    """

    def __init__(self, in_features, out_features, win_weight=2048, wout_weight=2048, dropout=0., bias=True):
        super(GraphConvolution, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        # self.weight = Parameter(torch.Tensor(in_features, out_features))
        self.weight = Parameter(torch.Tensor(win_weight, wout_weight))
        if bias:
            self.bias = Parameter(torch.Tensor(1, 1, wout_weight))
        else:
            self.register_parameter('bias', None)

        if dropout>0.01:
            self.drop_flag = True
            self.dropout_layer = nn.Dropout(p=dropout)
        else:
            self.drop_flag = False

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            #self.bias.data.uniform_(-stdv, stdv)
            nn.init.kaiming_normal_(self.weight, a=0)
            nn.init.constant_(self.bias, 0.0)

    def forward(self, input, adj):
        if self.drop_flag:
            input = self.dropout_layer(input)
        # print(input[:,0,:].shape)
        # print(self.weight.shape)
        support = torch.matmul(input, self.weight)
        output = torch.matmul(adj, support)

        if self.bias is not None:
            return output + self.bias
        else:
            return output

    def __repr__(self):
        return self.__class__.__name__ + ' (' \
               + str(self.in_features) + ' -> ' \
               + str(self.out_features) + ')'


def gen_adj_sim(part_features1,part_features2,enable_mask=False):
    """
    default similarity function, calculate similarity between two input features

    enable_mask: encourage sparse connections when similarities higher than the threshold

    Args:
        num_block: abandoned feature
        part_features1: vector1, size = b,c,n1
        part_features2: vector2, size = b,c,n2
        enable_mask: True/ False

    Returns:
        similarity matrices: size: n1 * n2

    """


    b,p,c = part_features1.size()
    features = part_features1
    #edge = torch.zeros(b,num_block, num_block)  # 33, 33
    feature2 = part_features2
    edge =torch.matmul(features, feature2)#.detach()
    #edge = torch.div(edge, float(p*p))
    #edge = torch.mul(torch.sign(edge), torch.sqrt(torch.abs(edge) + 1e-12))
    threshold = random.uniform(0.5, 0.6)

    zeros_vec = torch.zeros_like(edge)
    ones_vec = torch.ones_like(edge)
    eye_vec = torch.eye(p)


    value,idx = edge.max(dim=1)
    value_thre = value.unsqueeze(1)*threshold


    if enable_mask:
        mask = edge-value_thre.repeat(1,p,1)
        mask = torch.where(mask >= 0, ones_vec, zeros_vec)+eye_vec.cuda()
        return edge *mask
    #edge=F.softmax(edge,dim=2)
    #print(mask[0].sum(1))

    else:
        return edge

def norm_adj_batch(A):
    # print(A)
    # A = A.clamp(min=1e-11)
    D = torch.pow(A.sum(2).float(), -0.5)   # nam
    # print(torch.isnan(D).any())
    # raise '88888888888888'
    # D = torch.pow(A.sum(2).float() + 1e-10, -0.5)
    D = torch.diag_embed(D)
    adj = torch.matmul(A, D).permute(0,2,1)
    adj = torch.matmul(adj, D)
    return adj
class Embed(nn.Module):
    """Embedding module"""
    def __init__(self, dim_in=1024, dim_out=128, factor=2, convs=False):
        super(Embed, self).__init__()
        self.convs = convs
        if self.convs:
            self.transfer = nn.Sequential(
                nn.Conv2d(dim_in, dim_in//factor, kernel_size=1),
                nn.BatchNorm2d(dim_in//factor),
                nn.ReLU(inplace=True),
                nn.Conv2d(dim_in//factor, dim_in//factor, kernel_size=3, padding=1),
                nn.BatchNorm2d(dim_in//factor),
                nn.ReLU(inplace=True), 
                nn.Conv2d(dim_in//factor, dim_out, kernel_size=1),
                nn.BatchNorm2d(dim_out),
                nn.ReLU(inplace=True)              
            )
        else:
            self.transfer = nn.Sequential(
                nn.Conv2d(dim_in, dim_out, kernel_size=1),
                nn.BatchNorm2d(dim_out),
                nn.ReLU(inplace=True) 
            )


    def forward(self, x):
        x = self.transfer(x)
        return x

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)

class Presnet(nn.Module):
    def __init__(self, img_size=(224, 224),model_path="",embed_dim=96,depth_token_only=2, qkv_bias=True, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.5,stride_size=16,
                 norm_layer=nn.LayerNorm,num_stripes=2,last_stride=1,pre_trained=None,multi_head=0,**kwargs):
        super().__init__()
        # self.vit_base=vit_base_patch16_224_TransReID(img_size=img_size,stride_size=stride_size,model_path=model_path,**kwargs)
        self.base = ResNet(last_stride=last_stride,
                   block=Bottleneck,
                   layers=[3, 4, 6, 3])
        self.num_stripes = 4   #4    # num_stripes = 8

        self.num_features=2048
        self.norm1 = norm_layer(self.num_features)
        self.norm2 = norm_layer(self.num_features)
        # self.normt = norm_layer(self.num_features)
        self.multi_head=multi_head

        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.maxpool = nn.MaxPool1d(4)
        # self.reduction_layer = nn.Sequential(
        #     nn.BatchNorm2d(256),
        #     nn.ReLU(),
        #     nn.Conv2d(256, 64, kernel_size=1, padding=0, bias=True),
        # )
        # self.reduction_layer.apply(weights_init_kaiming)
        # self.reduction_layer2 = nn.Sequential(
        #     nn.BatchNorm2d(256),
        #     nn.ReLU(),
        #     nn.Conv2d(256, 64, kernel_size=1, padding=0, bias=True),
        # )
        # self.reduction_layer2.apply(weights_init_kaiming)

        self.apply(self._init_weights)
        if model_path and pre_trained=="imagenet":
            self.base.load_param(model_path)
            print("load_param  imagenet")
        self.backbone = nn.Sequential(
            self.base.conv1,
            self.base.bn1,
            self.base.relu,
            self.base.maxpool,
            # self.base.layer1,
            # self.base.layer2,
        )
        self.layer1 = self.base.layer1
        self.layer2 = self.base.layer2
        # self.layer3 = self.base.layer3     # copy.deepcopy(self.base.layer3)
        # self.layer2 = copy.deepcopy(self.base.layer2)
        self.layer3 = copy.deepcopy(self.base.layer3)
        self.layer4 = copy.deepcopy(self.base.layer4)
        # self.layer4 = copy.deepcopy(self.base.layer4)
        # self.layer2_2  = copy.deepcopy(self.base.layer2)
        self.layer3_2 = copy.deepcopy(self.base.layer3)
        self.layer4_2 = copy.deepcopy(self.base.layer4)
        # self.layer3_t = copy.deepcopy(self.base.layer3)
        # self.layer4_t = copy.deepcopy(self.base.layer4)        
        # self.pool_conv = nn.ModuleList()
        # for i in range(self.num_stripes):
        #     self.pool_conv.append(nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3))
        self.sigmoid = nn.Sigmoid()
        
        # self.channel_attens = nn.ModuleList()
        # self.linear_pool = nn.ModuleList()
        # for i in range(4):
            # self.linear_pool.append(nn.Sequential(
            # nn.Linear(256, 1),
            # nn.ReLU(inplace=True),
            # nn.Linear(64, 1), 
            # nn.Sigmoid()
            # ))
            # self.channel_attens.append(nn.Sequential(
            #     nn.Linear(512, 512// 8),
            #     nn.ReLU(inplace=True),
            #     nn.Linear(512 // 8, 512), 
            #     # nn.Sigmoid()
            # ))
        # channel_nums = 2048/self.num_stripes
        # { "Spatial Conv", "Avg", "Max", "Avg and Max"}
        self.spatial_mode = "Spatial Conv"
        # self.local_convs = nn.ModuleList()
        # # self.llhh = 2048    #  // 4
        # for i in range(self.num_stripes):
        #     self.local_convs.append(nn.Sequential(
        #         nn.Conv2d(512, 512//4, kernel_size = 1),
        #         nn.BatchNorm2d(512//4), 
        #         nn.ReLU(inplace=True),
        #         nn.Conv2d(512//4, 512, kernel_size = 1), 
        #     # nn.Sigmoid()
        #     ))
        self.distill_sssA = nn.ModuleList()
        for innini in range(5):
            self.distill_sssA.append(SpatialAttention())
        
        self.cls_tokens_list = []
        self.blocks_token_only = nn.ModuleList([
            LayerScale_Block_CA(
                dim=embed_dim*8, num_heads=8, mlp_ratio=4.0, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=0.0, attn_drop=0.0, drop_path=0.0, norm_layer=norm_layer,
                act_layer=nn.GELU,Attention_block='Class_Attention',
                Mlp_block='Mlp_old',init_values=1e-4,num_stripes=num_stripes)
            for i in range(1)])            
        # self.trans_embed = nn.ModuleList()
        # for i in range(4):
        #     self.trans_embed.append(Embed(512, 1024))
        # for i in range(4):
        #     self.trans_embed.append(Embed(512, 2048))

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, x,tdeep=1):
        # B = x[0].shape[0]    # x[0]
        B, _, h, w = x[0].shape
        x_b = self.backbone(x[0])   # x[0]

        x_1 = self.layer1(x_b)
        x_11 = self.layer2(x_1)
        global_feat_3 = self.layer3(x_11)

        global_feat = self.layer4(global_feat_3)
  
        feat_dim=global_feat.size(1)


        local_feat_3 = self.layer3_2(x_11)
        local_feat = self.layer4_2(local_feat_3)

        # local_feat_t = self.layer3_t(x_11)
        # local_feat_t = self.layer4_t(local_feat_t)
        
        H = local_feat.shape[-1]
  
          # Remove bg
        # if x[1] is not None:
        #     mask = x[1][:, 1:, :, :]
        #     B, N, HH, WW = mask.shape
        # else:
        #     B, _, HH, WW = x[0].shape
        #     N = 4
        #     mask = x[0].new_zeros(B, 4, HH, WW)
  
        # mask = F.interpolate(mask, global_feat.shape[2:])
        # local_feat_map = torch.mul(mask.unsqueeze(
        #     dim=2), global_feat.unsqueeze(dim=1))  # (B, N, C, h, w)
        # print(local_feat_map.shape)
        # raise "4444444"
        # local_feat_map = local_feat_map.view(B, -1, h, w)
        # local_feat_map = local_feat_map.reshape(B, -1, 16, 16)
        
        # vis_score = mask.sum(dim=[2, 3]) + 1  # Laplace平滑

        # local_feat_before = F.adaptive_avg_pool2d(local_feat_map, output_size=(1, 1)).view(B, N, 2048).permute(
        #     [0, 2, 1]) * (h * w / vis_score.unsqueeze(dim=1))  # (B, C, N)

  
        global_feat = global_feat.view(B,feat_dim,-1).transpose(1, 2).contiguous()    # [8, 196, 2048]
        local_feat = local_feat.view(B,feat_dim,-1).transpose(1, 2).contiguous()    # [8, 196, 2048]
        # local_feat_t = local_feat_t.view(B,feat_dim,-1).transpose(1, 2).contiguous()


        if tdeep==0:
            global_feat = self.avgpool(global_feat.transpose(1, 2))  # B C 1
            global_feat = torch.flatten(global_feat, 1)        
            return [global_feat]

        global_feat = self.norm1(global_feat).clamp(min=1e-11)
        local_feat = self.norm2(local_feat).clamp(min=1e-11)
        # global_feat = torch.relu(self.norm1(global_feat))
        # local_feat = torch.relu(self.norm2(local_feat))
        # local_feat_t = self.normt(local_feat_t)
        # if False:    #mask is None:
        #     local_feat_chunk = local_feat.chunk(self.num_stripes, dim=2)     # 按 channel 划分为 4 个部分  lo_fe_chu[0].shape == torch.Size([8, 196, 512])
        # elif SPATIAL_AND_CHANNEL:
            # 局部新划分
        local_split = "chunk"
        if local_split == "chunk":
            local_feat = local_feat.transpose(1, 2).view(B,feat_dim, H ,H).contiguous()
            local_feat_chunk = local_feat.chunk(self.num_stripes, dim=1)
            
            # local_feat_t = local_feat_t.transpose(1, 2).view(B,feat_dim, 16 ,16).contiguous()
            # local_feat_t_chunk = local_feat_t.chunk(2, dim=1)            
            
            global_feat_inter = global_feat.transpose(1, 2).view(B,feat_dim, 16 ,16).contiguous()
   
        #######  disitll spatial att
        # spatial_att_list = [self.distill_sssA[0](global_feat_inter)]
        # for ispai in range(self.num_stripes):
        #     spatial_att_list.append(self.distill_sssA[ispai+1](local_feat_chunk[ispai]))
        # global_feat = global_feat_inter + global_feat_inter*spatial_att_list[0]
        # tmep_feat = []
        # for sssssi in range(self.num_stripes):
        #     temp = local_feat_chunk[sssssi] + (spatial_att_list[sssssi+1] * local_feat_chunk[sssssi])
        #     tmep_feat.append(temp)
     
        # TODO 空间信息 图节点 聚类局部
        if len(self.cls_tokens_list) == 0:   # (B, 512, 256) => (B, 512, 257) 
            self.cls_tokens_list = [nn.Parameter(global_feat.view(B, 2048, -1).transpose(1, 2).contiguous().mean(dim=1, keepdim=True)).cuda()]
            # Use the mean features as the initial value of the global node
            for i in range(self.num_stripes):
                # print(local_feat_chunk[i].view(B, 512, -1).contiguous().shape)
                mean_features = local_feat_chunk[i].view(B, 512, -1).transpose(1, 2).contiguous().mean(dim=1, keepdim=True).cuda()
                self.cls_global_node = nn.Parameter(mean_features).cuda()
                self.cls_tokens_list.append(self.cls_global_node)
        else:
            pass
        ### test
        for i , blk in enumerate(self.blocks_token_only):
            cls_tokens = blk(global_feat.view(B, 2048, -1).transpose(1, 2).contiguous(), local_feat_chunk, self.cls_tokens_list)
        # print(self.training)
        # raise "777"
        if self.training:
            for i in range(5):
                # print(torch.isnan(cls_tokens[i]).any())
                new_global_node_feature = cls_tokens[i][:,0,:].unsqueeze(1)
                self.cls_tokens_list[i].data = new_global_node_feature.detach()  
                # local_spatial_node = local_feat.view(B, feat_dim, -1)
        # res_local_feat = []
        # for i in range(4):
        #     res_local_feat.append(self.local_convs[i](local_feat_chunk[i]))
        
        # res_local_feat  = self.local_convs(local_feat_chunk)
        # print(len(res_local_feat))
        global_feat = self.avgpool(global_feat.transpose(1, 2))  # B C 1
        global_feat = torch.flatten(global_feat, 1)
        # global_feat = self.avgpool(global_feat.view(B, 2048, -1).transpose(1, 2).contiguous().transpose(1, 2))  # B C 1
        # global_feat = torch.flatten(global_feat, 1)
        
        final_feats=[global_feat]
        if SPATIAL_AND_CHANNEL:
            B, C, h, w = local_feat.shape

                
            # spatial_atts = []

            # spatials_ca = 'torch.cat(spatial_atts, dim=1)'

            for inni in range(self.num_stripes):
                  # 

                # spatial_atts.append(spatial_att)
                # x1 = self.maxpool(temp_chunk.view(B, 2048//self.num_stripes, -1).contiguous())
                # x1 = self.avgpool(x1)

                # x1 = self.avgpool(temp_chunk.view(B, 2048//self.num_stripes, -1).contiguous())  # B C 1
                x1 = self.avgpool(local_feat_chunk[inni].view(B, 512, -1).contiguous())  # B C 1
                x1 = torch.flatten(x1, 1)
                final_feats.append(x1)

            # for innnni in range(2):
            # # spatial_atts.append(spatial_att)
            #     x1_t = self.avgpool(local_feat_t_chunk[innnni].view(B, 1024, -1).contiguous())  # B C 1
            #     x1_t = torch.flatten(x1_t, 1)
            #     final_feats.append(x1_t)

            # tea_inter.append(local_feat_t_chunk[0])
            # tea_inter.append(local_feat_t_chunk[1])


            # embed_1024 = []
            # embed_2048 = []
            # for iiii in range(4):
            #     embed_1024.append(self.trans_embed[iiii](local_feat_chunk[iiii]))
            # for iiii in range(4):
            #     embed_2048.append(self.trans_embed[-(iiii+1)](local_feat_chunk[iiii]))

        return_feats=[]
        for iaa in range(self.num_stripes+1):             # self.num_stripes+  +2
            
            ### test
            tmp_feat=final_feats[iaa]+cls_tokens[iaa][:,0,:]

            # tmp_feat = res[:,i,:]
            # print(tmp_feat.size())
            return_feats.append(tmp_feat)

        # spatials_ca 没有用到
        # if SPATIAL_AND_CHANNEL:
        # print(len(return_feats))
        # print(return_feats[1].shape)
        # raise "ss"
        return return_feats  # loss44 embed_1024+embed_2048,




